package class_0822Graph;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * @author AutumnWhisper
 */
public class Coding_Prim {
    public static class EdgeComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;  // 按权重从小到大排序
        }
    }

    public static Set<Edge> primMST(Graph graph) {
        PriorityQueue<Edge> pq = new PriorityQueue<>(new EdgeComparator());  // 优先队列，存储当前可用的边
        HashSet<Node> set = new HashSet<>();  // 用于跟踪已经加入MST的节点
        Set<Edge> result = new HashSet<>();  // 用于存储MST中的边

        for (Node node : graph.nodes.values()) {  // 遍历图中的所有节点
            if (!set.contains(node)) {  // 如果节点尚未包含在MST中
                set.add(node);  // 将节点加入MST集合
                for (Edge edge : node.edges) {  // 将该节点的所有边加入优先队列
                    pq.offer(edge);
                }

                while (!pq.isEmpty()) {  // 不断从优先队列中选择权重最小的边
                    Edge edge = pq.poll();  // 取出权重最小的边
                    Node toNode = edge.to;  // 获得这条边指向的节点

                    if (!set.contains(toNode)) {  // 如果该节点不在MST中
                        set.add(toNode);  // 将节点加入MST集合
                        result.add(edge);  // 将这条边加入结果集中
                        for (Edge nextEdge : toNode.edges) {  // 将该节点的所有边加入优先队列
                            if (!pq.contains(nextEdge)) {  // 防止重复添加
                                pq.offer(nextEdge);
                            }
                        }
                    }
                }
            }
        }
        return result;  // 返回构建好的MST
    }
}
